{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/vodka/anaconda3/lib/python3.7/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:541: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint8 = np.dtype([(\"qint8\", np.int8, 1)])\n",
      "/home/vodka/anaconda3/lib/python3.7/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:542: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint8 = np.dtype([(\"quint8\", np.uint8, 1)])\n",
      "/home/vodka/anaconda3/lib/python3.7/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:543: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint16 = np.dtype([(\"qint16\", np.int16, 1)])\n",
      "/home/vodka/anaconda3/lib/python3.7/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:544: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint16 = np.dtype([(\"quint16\", np.uint16, 1)])\n",
      "/home/vodka/anaconda3/lib/python3.7/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:545: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint32 = np.dtype([(\"qint32\", np.int32, 1)])\n",
      "/home/vodka/anaconda3/lib/python3.7/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:550: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  np_resource = np.dtype([(\"resource\", np.ubyte, 1)])\n",
      "/home/vodka/anaconda3/lib/python3.7/site-packages/tensorflow/python/framework/dtypes.py:516: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint8 = np.dtype([(\"qint8\", np.int8, 1)])\n",
      "/home/vodka/anaconda3/lib/python3.7/site-packages/tensorflow/python/framework/dtypes.py:517: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint8 = np.dtype([(\"quint8\", np.uint8, 1)])\n",
      "/home/vodka/anaconda3/lib/python3.7/site-packages/tensorflow/python/framework/dtypes.py:518: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint16 = np.dtype([(\"qint16\", np.int16, 1)])\n",
      "/home/vodka/anaconda3/lib/python3.7/site-packages/tensorflow/python/framework/dtypes.py:519: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint16 = np.dtype([(\"quint16\", np.uint16, 1)])\n",
      "/home/vodka/anaconda3/lib/python3.7/site-packages/tensorflow/python/framework/dtypes.py:520: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint32 = np.dtype([(\"qint32\", np.int32, 1)])\n",
      "/home/vodka/anaconda3/lib/python3.7/site-packages/tensorflow/python/framework/dtypes.py:525: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  np_resource = np.dtype([(\"resource\", np.ubyte, 1)])\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import json\n",
    "import logging\n",
    "import random\n",
    "from collections import defaultdict\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import pytorch_lightning as pl\n",
    "from tqdm import tqdm\n",
    "from sklearn.metrics import accuracy_score\n",
    "from transformers import (\n",
    "    DataProcessor,\n",
    "    InputExample,\n",
    "    BertConfig,\n",
    "    BertTokenizer,\n",
    "    BertForSequenceClassification,\n",
    "    glue_convert_examples_to_features,\n",
    ")\n",
    "\n",
    "DEVICE = torch.device('cuda:1' if torch.cuda.is_available else 'cpu')\n",
    "\n",
    "# 预训练模型路径\n",
    "PRETRAINED_PATH = 'hfl/chinese_roberta_wwm_ext_pytorch/'\n",
    "# 实体链接训练路径\n",
    "EL_SAVE_PATH = './pytorch-lightning-checkpoints/EntityLinking/'\n",
    "# 实体类别推断训练路径\n",
    "ET_SAVE_PATH = './pytorch-lightning-checkpoints/EntityTyping/'\n",
    "\n",
    "# 项目数据路径\n",
    "DATA_PATH = './data/'\n",
    "\n",
    "# CCKS2020实体链指竞赛原始路径\n",
    "RAW_PATH = DATA_PATH + 'ccks2020_el_data_v1/'\n",
    "\n",
    "# 预处理后导出的pickle文件路径\n",
    "PICKLE_PATH = DATA_PATH + 'pickle/'\n",
    "if not os.path.exists(PICKLE_PATH):\n",
    "    os.mkdir(PICKLE_PATH)\n",
    "\n",
    "# 预测结果的文件路径\n",
    "RESULT_PATH = DATA_PATH + 'result/'\n",
    "if not os.path.exists(RESULT_PATH):\n",
    "    os.mkdir(RESULT_PATH)\n",
    "\n",
    "# 训练、验证、推断所需的tsv文件路径\n",
    "TSV_PATH = DATA_PATH + 'tsv/'\n",
    "if not os.path.exists(TSV_PATH):\n",
    "    os.mkdir(TSV_PATH)\n",
    "\n",
    "# 训练结果的CheckPoint文件路径\n",
    "CKPT_PATH = './ckpt/'\n",
    "\n",
    "PICKLE_DATA = {\n",
    "    # 实体名称对应的KBID列表\n",
    "    'ENTITY_TO_KBIDS': None,\n",
    "    # KBID对应的实体名称列表\n",
    "    'KBID_TO_ENTITIES': None,\n",
    "    # KBID对应的属性文本\n",
    "    'KBID_TO_TEXT': None,\n",
    "    # KBID对应的实体类型列表（注意：一个实体可能对应'|'分割的多个类型）\n",
    "    'KBID_TO_TYPES': None,\n",
    "    # KBID对应的属性列表\n",
    "    'KBID_TO_PREDICATES': None,\n",
    "\n",
    "    # 索引类型映射列表\n",
    "    'IDX_TO_TYPE': None,\n",
    "    # 类型索引映射字典\n",
    "    'TYPE_TO_IDX': None,\n",
    "}\n",
    "\n",
    "\n",
    "def read_pickle():\n",
    "    for k in PICKLE_DATA:\n",
    "        filename = k + '.pkl'\n",
    "        if os.path.exists(PICKLE_PATH + filename):\n",
    "            PICKLE_DATA[k] = pd.read_pickle(PICKLE_PATH + filename)\n",
    "        else:\n",
    "            print(f'File {filename} not Exist!')\n",
    "\n",
    "\n",
    "def set_random_seed(seed):\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed(seed)\n",
    "    torch.cuda.manual_seed_all(seed)\n",
    "    torch.backends.cudnn.benchmark = False\n",
    "    torch.backends.cudnn.deterministic = True\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "\n",
    "class DataFramePreprocessor:\n",
    "    \"\"\"生成模型训练、验证、推断所需的tsv文件\"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "        pass\n",
    "\n",
    "    def process_link_data(self, input_path, output_path, max_negs=-1):\n",
    "\n",
    "        entity_to_kbids = PICKLE_DATA['ENTITY_TO_KBIDS']\n",
    "        # print(\"entity_to_kbids\")\n",
    "        kbid_to_text = PICKLE_DATA['KBID_TO_TEXT']\n",
    "        # print(kbid_to_text)\n",
    "        kbid_to_predicates = PICKLE_DATA['KBID_TO_PREDICATES']\n",
    "        link_dict = defaultdict(list)\n",
    "\n",
    "        with open(input_path, 'r', encoding='utf-8') as f:\n",
    "            for line in tqdm(f):\n",
    "                line = json.loads(line)\n",
    "\n",
    "                for data in line['mention_data']:\n",
    "                    # 对测试集特殊处理\n",
    "                    if 'kb_id' not in data:\n",
    "                        data['kb_id'] = '0'\n",
    "\n",
    "                    # KB中不存在的实体不进行链接\n",
    "                    if not data['kb_id'].isdigit():\n",
    "                        continue\n",
    "\n",
    "                    entity = data['mention']\n",
    "                    kbids = list(entity_to_kbids[entity])\n",
    "                    random.shuffle(kbids)\n",
    "\n",
    "                    num_negs = 0\n",
    "                    for kbid in kbids:\n",
    "                        if num_negs >= max_negs > 0 and kbid != data['kb_id']:\n",
    "                            continue\n",
    "\n",
    "                        link_dict['text_id'].append(line['text_id'])\n",
    "                        link_dict['entity'].append(entity)\n",
    "                        link_dict['offset'].append(data['offset'])\n",
    "                        link_dict['short_text'].append(line['text'])\n",
    "                        link_dict['kb_id'].append(kbid)\n",
    "                        link_dict['kb_text'].append(kbid_to_text[kbid])\n",
    "                        link_dict['kb_predicate_num'].append(len(kbid_to_predicates[kbid]))\n",
    "                        if kbid != data['kb_id']:\n",
    "                            link_dict['predict'].append(0)\n",
    "                            num_negs += 1\n",
    "                        else:\n",
    "                            link_dict['predict'].append(1)\n",
    "\n",
    "        link_data = pd.DataFrame(link_dict)\n",
    "        link_data.to_csv(output_path, index=False, sep='\\t')\n",
    "\n",
    "    def process_type_data(self, input_path, output_path):\n",
    "        kbid_to_types = PICKLE_DATA['KBID_TO_TYPES']\n",
    "        type_dict = defaultdict(list)\n",
    "\n",
    "        with open(input_path, 'r', encoding='utf-8') as f:\n",
    "            for line in tqdm(f):\n",
    "                line = json.loads(line)\n",
    "\n",
    "                for data in line['mention_data']:\n",
    "                    entity = data['mention']\n",
    "\n",
    "                    # 测试集特殊处理\n",
    "                    if 'kb_id' not in data:\n",
    "                        entity_type = ['Other']\n",
    "                    elif data['kb_id'].isdigit():\n",
    "                        entity_type = kbid_to_types[data['kb_id']]\n",
    "                    else:\n",
    "                        entity_type = data['kb_id'].split('|')\n",
    "                        for x in range(len(entity_type)):\n",
    "                            entity_type[x] = entity_type[x][4:]\n",
    "                    for e in entity_type:\n",
    "                        type_dict['text_id'].append(line['text_id'])\n",
    "                        type_dict['entity'].append(entity)\n",
    "                        type_dict['offset'].append(data['offset'])\n",
    "                        type_dict['short_text'].append(line['text'])\n",
    "                        type_dict['type'].append(e)\n",
    "\n",
    "        type_data = pd.DataFrame(type_dict)\n",
    "        type_data.to_csv(output_path, index=False, sep='\\t')\n",
    "\n",
    "    def run(self):\n",
    "        self.process_link_data(\n",
    "            input_path=RAW_PATH + 'train.json',\n",
    "            output_path=TSV_PATH + 'EL_TRAIN.tsv',\n",
    "            max_negs=2,\n",
    "        )\n",
    "        logger.info('Process EL_TRAIN Finish.')\n",
    "        self.process_link_data(\n",
    "            input_path=RAW_PATH + 'dev.json',\n",
    "            output_path=TSV_PATH + 'EL_VALID.tsv',\n",
    "            max_negs=-1,\n",
    "        )\n",
    "        logger.info('Process EL_VALID Finish.')\n",
    "        self.process_link_data(\n",
    "            input_path=RAW_PATH + 'test.json',\n",
    "            output_path=TSV_PATH + 'EL_TEST.tsv',\n",
    "            max_negs=-1,\n",
    "        )\n",
    "        logger.info('Process EL_TEST Finish.')\n",
    "\n",
    "        self.process_type_data(\n",
    "            input_path=RAW_PATH + 'train.json',\n",
    "            output_path=TSV_PATH + 'ET_TRAIN.tsv',\n",
    "        )\n",
    "        logger.info('Process ET_TRAIN Finish.')\n",
    "        self.process_type_data(\n",
    "            input_path=RAW_PATH + 'dev.json',\n",
    "            output_path=TSV_PATH + 'ET_VALID.tsv',\n",
    "        )\n",
    "        logger.info('Process ET_VALID Finish.')\n",
    "        self.process_type_data(\n",
    "            input_path=RAW_PATH + 'test.json',\n",
    "            output_path=TSV_PATH + 'ET_TEST.tsv',\n",
    "        )\n",
    "        logger.info('Process ET_TEST Finish.')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "\n",
    "class PicklePreprocessor:\n",
    "    \"\"\"生成全局变量Pickle文件的预处理器\"\"\"\n",
    "\n",
    "    def __init__(self):\n",
    "\n",
    "        # 实体名称对应的KBID列表  {\"张健\"} -> \"10001\"\n",
    "        self.entity_to_kbids = defaultdict(set)\n",
    "\n",
    "        # KBID对应的实体名称列表 \"10001\" -> {\"张健\"}\n",
    "        self.kbid_to_entities = dict()\n",
    "\n",
    "        # KBID对应的属性文本  \"10001\" -> {\"政治面貌:中共党员\",\"义项描述:潜山县塔畈乡副主任科员、纪委副书记\",\"性别:男\",\n",
    "        # \"学历:大专\",\"中文名:张健\"}\n",
    "        self.kbid_to_text = dict()\n",
    "\n",
    "        # KBID对应的实体类型列表 \"10001\" -> {\"Person\"}\n",
    "        self.kbid_to_types = dict()\n",
    "\n",
    "        # KBID对应的属性列表 \"10001\" -> {\"政治面貌\",\"义项描述\",\"性别\",\"学历\",\"中文名\"}\n",
    "        self.kbid_to_predicates = dict()\n",
    "\n",
    "        # 索引类型映射列表 [\"Person\"]\n",
    "        self.idx_to_type = list()\n",
    "\n",
    "        # 类型索引映射字典 {\"Person\":0}\n",
    "        self.type_to_idx = dict()\n",
    "\n",
    "    def run(self, shuffle_text=True):\n",
    "        with open(RAW_PATH + 'kb.json', 'r', encoding='utf-8') as f:\n",
    "            for line in tqdm(f):\n",
    "                line = json.loads(line)\n",
    "\n",
    "                kbid = line['subject_id']\n",
    "                # 将实体名与别名合并\n",
    "                entities = set(line['alias'])\n",
    "                entities.add(line['subject'])\n",
    "                for entity in entities:\n",
    "                    self.entity_to_kbids[entity].add(kbid)\n",
    "                self.kbid_to_entities[kbid] = entities\n",
    "\n",
    "                text_list, predicate_list = [], []\n",
    "                for x in line['data']:\n",
    "                    # 简单拼接predicate与object，这部分可以考虑别的方法尝试\n",
    "                    text_list.append(':'.join([x['predicate'].strip(), x['object'].strip()]))\n",
    "                    predicate_list.append(x['predicate'].strip())\n",
    "                if shuffle_text:  # 对属性文本随机打乱顺序\n",
    "                    random.shuffle(text_list)\n",
    "                self.kbid_to_predicates[kbid] = predicate_list\n",
    "                self.kbid_to_text[kbid] = ' '.join(text_list)\n",
    "\n",
    "                # 删除文本中的特殊字符\n",
    "                for c in ['\\r', '\\t', '\\n']:\n",
    "                    self.kbid_to_text[kbid] = self.kbid_to_text[kbid].replace(c, '')\n",
    "\n",
    "                type_list = line['type'].split('|')\n",
    "                self.kbid_to_types[kbid] = type_list\n",
    "                for t in type_list:\n",
    "                    if t not in self.type_to_idx:\n",
    "                        self.type_to_idx[t] = len(self.idx_to_type)\n",
    "                        self.idx_to_type.append(t)\n",
    "\n",
    "        # 保存pickle文件\n",
    "        pd.to_pickle(self.entity_to_kbids, PICKLE_PATH + 'ENTITY_TO_KBIDS.pkl')\n",
    "        pd.to_pickle(self.kbid_to_entities, PICKLE_PATH + 'KBID_TO_ENTITIES.pkl')\n",
    "        pd.to_pickle(self.kbid_to_text, PICKLE_PATH + 'KBID_TO_TEXT.pkl')\n",
    "        pd.to_pickle(self.kbid_to_types, PICKLE_PATH + 'KBID_TO_TYPES.pkl')\n",
    "        pd.to_pickle(self.kbid_to_predicates, PICKLE_PATH + 'KBID_TO_PREDICATES.pkl')\n",
    "        pd.to_pickle(self.idx_to_type, PICKLE_PATH + 'IDX_TO_TYPE.pkl')\n",
    "        pd.to_pickle(self.type_to_idx, PICKLE_PATH + 'TYPE_TO_IDX.pkl')\n",
    "        logger.info('Process Pickle File Finish.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EntityLinkingProcessor(DataProcessor):\n",
    "    \"\"\"实体链接数据处理\"\"\"\n",
    "\n",
    "    def get_train_examples(self, file_path):\n",
    "        return self._create_examples(\n",
    "            self._read_tsv(file_path),\n",
    "            set_type='train',\n",
    "        )\n",
    "\n",
    "    def get_dev_examples(self, file_path):\n",
    "        return self._create_examples(\n",
    "            self._read_tsv(file_path),\n",
    "            set_type='valid',\n",
    "        )\n",
    "\n",
    "    def get_test_examples(self, file_path):\n",
    "        return self._create_examples(\n",
    "            self._read_tsv(file_path),\n",
    "            set_type='test',\n",
    "        )\n",
    "\n",
    "    def get_labels(self):\n",
    "        return ['0', '1']\n",
    "\n",
    "    def _create_examples(self, lines, set_type):\n",
    "        examples = []\n",
    "        for i, line in enumerate(lines):\n",
    "            if i == 0:\n",
    "                continue\n",
    "            guid = f'{set_type}-{i}'\n",
    "            text_a = line[1] + ' ' + line[3]\n",
    "            text_b = line[5]\n",
    "            label = line[-1]\n",
    "            examples.append(InputExample(\n",
    "                guid=guid,\n",
    "                text_a=text_a,\n",
    "                text_b=text_b,\n",
    "                label=label,\n",
    "            ))\n",
    "        return examples\n",
    "\n",
    "    def create_dataloader(self, examples, tokenizer, max_length=384,\n",
    "                          shuffle=False, batch_size=32, use_pickle=False):\n",
    "        pickle_name = 'EL_FEATURE_' + examples[0].guid.split('-')[0].upper() + '.pkl'\n",
    "        if use_pickle:\n",
    "            features = pd.read_pickle(PICKLE_PATH + pickle_name)\n",
    "        else:\n",
    "            features = glue_convert_examples_to_features(\n",
    "                examples,\n",
    "                tokenizer,\n",
    "                label_list=self.get_labels(),\n",
    "                max_length=max_length,\n",
    "                output_mode='classification',\n",
    "            )\n",
    "\n",
    "            pd.to_pickle(features, PICKLE_PATH + pickle_name)\n",
    "\n",
    "        dataset = torch.utils.data.TensorDataset(\n",
    "            torch.LongTensor([f.input_ids for f in features]),\n",
    "            torch.LongTensor([f.attention_mask for f in features]),\n",
    "            torch.LongTensor([f.token_type_ids for f in features]),\n",
    "            torch.LongTensor([f.label for f in features]),\n",
    "        )\n",
    "\n",
    "        dataloader = torch.utils.data.DataLoader(\n",
    "            dataset,\n",
    "            shuffle=shuffle,\n",
    "            batch_size=batch_size,\n",
    "            num_workers=2,\n",
    "        )\n",
    "        return dataloader\n",
    "\n",
    "    def generate_feature_pickle(self, max_length):\n",
    "        tokenizer = BertTokenizer.from_pretrained(\"hfl/chinese-roberta-wwm-ext\")\n",
    "\n",
    "        train_examples = self.get_train_examples(TSV_PATH + 'EL_TRAIN.tsv')\n",
    "        valid_examples = self.get_dev_examples(TSV_PATH + 'EL_VALID.tsv')\n",
    "        test_examples = self.get_test_examples(TSV_PATH + 'EL_TEST.tsv')\n",
    "\n",
    "        self.create_dataloader(\n",
    "            examples=train_examples,\n",
    "            tokenizer=tokenizer,\n",
    "            max_length=max_length,\n",
    "            shuffle=True,\n",
    "            batch_size=32,\n",
    "            use_pickle=False,\n",
    "        )\n",
    "        self.create_dataloader(\n",
    "            examples=valid_examples,\n",
    "            tokenizer=tokenizer,\n",
    "            max_length=max_length,\n",
    "            shuffle=False,\n",
    "            batch_size=32,\n",
    "            use_pickle=False,\n",
    "        )\n",
    "        self.create_dataloader(\n",
    "            examples=test_examples,\n",
    "            tokenizer=tokenizer,\n",
    "            max_length=max_length,\n",
    "            shuffle=False,\n",
    "            batch_size=32,\n",
    "            use_pickle=False,\n",
    "        )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EntityLinkingModel(pl.LightningModule):\n",
    "    \"\"\"实体链接模型\"\"\"\n",
    "\n",
    "    def __init__(self, max_length=384, batch_size=32, use_pickle=True):\n",
    "        super(EntityLinkingModel, self).__init__()\n",
    "        # 输入最大长度\n",
    "        self.max_length = max_length\n",
    "        self.batch_size = batch_size\n",
    "        self.use_pickle = use_pickle\n",
    "\n",
    "        self.tokenizer = BertTokenizer.from_pretrained(\"hfl/chinese-roberta-wwm-ext\")\n",
    "\n",
    "        self.bert = BertForSequenceClassification.from_pretrained(\n",
    "            \"hfl/chinese-roberta-wwm-ext\",\n",
    "            num_labels=1,\n",
    "        )\n",
    "\n",
    "        # 二分类损失函数\n",
    "        self.criterion = nn.BCEWithLogitsLoss()\n",
    "\n",
    "    def forward(self, input_ids, attention_mask, token_type_ids):\n",
    "        logits = self.bert(\n",
    "            input_ids=input_ids,\n",
    "            attention_mask=attention_mask,\n",
    "            token_type_ids=token_type_ids,\n",
    "        )[0]\n",
    "        return logits.squeeze()\n",
    "\n",
    "    def prepare_data(self):\n",
    "        self.processor = EntityLinkingProcessor()\n",
    "        self.train_examples = self.processor.get_train_examples(TSV_PATH + 'EL_TRAIN.tsv')\n",
    "        self.valid_examples = self.processor.get_dev_examples(TSV_PATH + 'EL_VALID.tsv')\n",
    "        self.test_examples = self.processor.get_test_examples(TSV_PATH + 'EL_TEST.tsv')\n",
    "\n",
    "        self.train_loader = self.processor.create_dataloader(\n",
    "            examples=self.train_examples,\n",
    "            tokenizer=self.tokenizer,\n",
    "            max_length=self.max_length,\n",
    "            shuffle=True,\n",
    "            batch_size=self.batch_size,\n",
    "            use_pickle=self.use_pickle,\n",
    "        )\n",
    "        self.valid_loader = self.processor.create_dataloader(\n",
    "            examples=self.valid_examples,\n",
    "            tokenizer=self.tokenizer,\n",
    "            max_length=self.max_length,\n",
    "            shuffle=False,\n",
    "            batch_size=self.batch_size,\n",
    "            use_pickle=self.use_pickle,\n",
    "        )\n",
    "        self.test_loader = self.processor.create_dataloader(\n",
    "            examples=self.test_examples,\n",
    "            tokenizer=self.tokenizer,\n",
    "            max_length=self.max_length,\n",
    "            shuffle=False,\n",
    "            batch_size=self.batch_size,\n",
    "            use_pickle=self.use_pickle,\n",
    "        )\n",
    "        print(\"finish\")\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        input_ids, attention_mask, token_type_ids, labels = batch\n",
    "        logits = self(input_ids, attention_mask, token_type_ids)\n",
    "        loss = self.criterion(logits, labels.float())\n",
    "\n",
    "        preds = (logits > 0).int()\n",
    "        acc = (preds == labels).float().mean()\n",
    "\n",
    "        tensorboard_logs = {'train_loss': loss, 'train_acc': acc}\n",
    "        return {'loss': loss, 'log': tensorboard_logs, 'progress_bar': tensorboard_logs}\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        input_ids, attention_mask, token_type_ids, labels = batch\n",
    "        logits = self(input_ids, attention_mask, token_type_ids)\n",
    "        loss = self.criterion(logits, labels.float())\n",
    "\n",
    "        preds = (logits > 0).int()\n",
    "        acc = (preds == labels).float().mean()\n",
    "\n",
    "        return {'val_loss': loss, 'val_acc': acc}\n",
    "\n",
    "    def validation_epoch_end(self, outputs):\n",
    "        val_loss = torch.stack([x['val_loss'] for x in outputs]).mean()\n",
    "        val_acc = torch.stack([x['val_acc'] for x in outputs]).mean()\n",
    "\n",
    "        tensorboard_logs = {'val_loss': val_loss, 'val_acc': val_acc}\n",
    "        return {'val_loss': val_loss, 'log': tensorboard_logs, 'progress_bar': tensorboard_logs}\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        return torch.optim.Adam([p for p in self.parameters() if p.requires_grad], lr=2e-5, eps=1e-8)\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return self.train_loader\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        return self.valid_loader\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EntityLinkingPredictor:\n",
    "\n",
    "    def __init__(self, ckpt_name, batch_size=8, use_pickle=True):\n",
    "        self.ckpt_name = ckpt_name\n",
    "        self.batch_size = batch_size\n",
    "        self.use_pickle = use_pickle\n",
    "\n",
    "    def generate_tsv_result(self, tsv_name, tsv_type='Valid'):\n",
    "        processor = EntityLinkingProcessor()\n",
    "        tokenizer = BertTokenizer.from_pretrained(\"hfl/chinese-roberta-wwm-ext\")\n",
    "\n",
    "        if tsv_type == 'Valid':\n",
    "            examples = processor.get_dev_examples(TSV_PATH + tsv_name)\n",
    "        elif tsv_type == 'Test':\n",
    "            examples = processor.get_test_examples(TSV_PATH + tsv_name)\n",
    "        else:\n",
    "            raise ValueError('tsv_type error')\n",
    "        dataloader = processor.create_dataloader(\n",
    "            examples=examples,\n",
    "            tokenizer=tokenizer,\n",
    "            max_length=384,\n",
    "            shuffle=False,\n",
    "            batch_size=self.batch_size,\n",
    "            use_pickle=self.use_pickle,\n",
    "        )\n",
    "\n",
    "        model = EntityLinkingModel.load_from_checkpoint(\n",
    "            checkpoint_path=CKPT_PATH + self.ckpt_name,\n",
    "        )\n",
    "        model.to(DEVICE)\n",
    "        model = nn.DataParallel(model)\n",
    "        model.eval()\n",
    "\n",
    "        result_list, logit_list = [], []\n",
    "        for batch in tqdm(dataloader):\n",
    "            for i in range(len(batch)):\n",
    "                batch[i] = batch[i].to(DEVICE)\n",
    "\n",
    "            input_ids, attention_mask, token_type_ids, labels = batch\n",
    "            logits = model(input_ids, attention_mask, token_type_ids)\n",
    "            preds = (logits > 0).int()\n",
    "\n",
    "            result_list.extend(preds.tolist())\n",
    "            logit_list.extend(logits.tolist())\n",
    "\n",
    "        tsv_data = pd.read_csv(TSV_PATH + tsv_name, sep='\\t')\n",
    "        tsv_data['logits'] = logit_list\n",
    "        tsv_data['result'] = result_list\n",
    "        result_name = tsv_name.split('.')[0] + '_RESULT.tsv'\n",
    "        tsv_data.to_csv(RESULT_PATH + result_name, index=False, sep='\\t')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EntityTypingProcessor(DataProcessor):\n",
    "    \"\"\"实体链接数据处理\"\"\"\n",
    "\n",
    "    def get_train_examples(self, file_path):\n",
    "        return self._create_examples(\n",
    "            self._read_tsv(file_path),\n",
    "            set_type='train',\n",
    "        )\n",
    "\n",
    "    def get_dev_examples(self, file_path):\n",
    "        return self._create_examples(\n",
    "            self._read_tsv(file_path),\n",
    "            set_type='valid',\n",
    "        )\n",
    "\n",
    "    def get_test_examples(self, file_path):\n",
    "        return self._create_examples(\n",
    "            self._read_tsv(file_path),\n",
    "            set_type='test',\n",
    "        )\n",
    "\n",
    "    def get_labels(self):\n",
    "        return PICKLE_DATA['IDX_TO_TYPE']\n",
    "\n",
    "    def _create_examples(self, lines, set_type):\n",
    "        examples = []\n",
    "        for i, line in enumerate(lines):\n",
    "            if i == 0:\n",
    "                continue\n",
    "            guid = f'{set_type}-{i}'\n",
    "            text_a = line[1]\n",
    "            text_b = line[3]\n",
    "            label = line[-1]\n",
    "            examples.append(InputExample(\n",
    "                guid=guid,\n",
    "                text_a=text_a,\n",
    "                text_b=text_b,\n",
    "                label=label,\n",
    "            ))\n",
    "        return examples\n",
    "\n",
    "    def create_dataloader(self, examples, tokenizer, max_length=64,\n",
    "                          shuffle=False, batch_size=64, use_pickle=False):\n",
    "        pickle_name = 'ET_FEATURE_' + examples[0].guid.split('-')[0].upper() + '.pkl'\n",
    "        if use_pickle:\n",
    "            features = pd.read_pickle(PICKLE_PATH + pickle_name)\n",
    "        else:\n",
    "            features = glue_convert_examples_to_features(\n",
    "                examples,\n",
    "                tokenizer,\n",
    "                label_list=self.get_labels(),\n",
    "                max_length=max_length,\n",
    "                output_mode='classification',\n",
    "            )\n",
    "            pd.to_pickle(features, PICKLE_PATH + pickle_name)\n",
    "\n",
    "        dataset = torch.utils.data.TensorDataset(\n",
    "            torch.LongTensor([f.input_ids for f in features]),\n",
    "            torch.LongTensor([f.attention_mask for f in features]),\n",
    "            torch.LongTensor([f.token_type_ids for f in features]),\n",
    "            torch.LongTensor([f.label for f in features]),\n",
    "        )\n",
    "\n",
    "        dataloader = torch.utils.data.DataLoader(\n",
    "            dataset,\n",
    "            shuffle=shuffle,\n",
    "            batch_size=batch_size,\n",
    "            num_workers=2,\n",
    "        )\n",
    "        return dataloader\n",
    "\n",
    "    def generate_feature_pickle(self, max_length):\n",
    "        tokenizer = BertTokenizer.from_pretrained(\"hfl/chinese-roberta-wwm-ext\")\n",
    "\n",
    "        train_examples = self.get_train_examples(TSV_PATH + 'ET_TRAIN.tsv')\n",
    "        valid_examples = self.get_dev_examples(TSV_PATH + 'ET_VALID.tsv')\n",
    "        test_examples = self.get_test_examples(TSV_PATH + 'ET_TEST.tsv')\n",
    "\n",
    "        self.create_dataloader(\n",
    "            examples=train_examples,\n",
    "            tokenizer=tokenizer,\n",
    "            max_length=max_length,\n",
    "            shuffle=True,\n",
    "            batch_size=32,\n",
    "            use_pickle=False,\n",
    "        )\n",
    "        self.create_dataloader(\n",
    "            examples=valid_examples,\n",
    "            tokenizer=tokenizer,\n",
    "            max_length=max_length,\n",
    "            shuffle=False,\n",
    "            batch_size=32,\n",
    "            use_pickle=False,\n",
    "        )\n",
    "        self.create_dataloader(\n",
    "            examples=test_examples,\n",
    "            tokenizer=tokenizer,\n",
    "            max_length=max_length,\n",
    "            shuffle=False,\n",
    "            batch_size=32,\n",
    "            use_pickle=False,\n",
    "        )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EntityTypingModel(pl.LightningModule):\n",
    "    \"\"\"实体类型推断模型\"\"\"\n",
    "\n",
    "    def __init__(self, max_length=64, batch_size=64, use_pickle=True):\n",
    "        super(EntityTypingModel, self).__init__()\n",
    "        # 输入最大长度\n",
    "        self.max_length = max_length\n",
    "        self.batch_size = batch_size\n",
    "        self.use_pickle = use_pickle\n",
    "        # 二分类损失函数\n",
    "        self.criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "        self.tokenizer = BertTokenizer.from_pretrained(\"hfl/chinese-roberta-wwm-ext\")\n",
    "\n",
    "        # 预训练模型\n",
    "        self.bert = BertForSequenceClassification.from_pretrained(\n",
    "            \"hfl/chinese-roberta-wwm-ext\",\n",
    "            num_labels=len(PICKLE_DATA['IDX_TO_TYPE']),\n",
    "        )\n",
    "\n",
    "    def forward(self, input_ids, attention_mask, token_type_ids):\n",
    "        return self.bert(\n",
    "            input_ids=input_ids,\n",
    "            attention_mask=attention_mask,\n",
    "            token_type_ids=token_type_ids,\n",
    "        )[0]\n",
    "\n",
    "    def prepare_data(self):\n",
    "        self.processor = EntityTypingProcessor()\n",
    "        self.train_examples = self.processor.get_train_examples(TSV_PATH + 'ET_TRAIN.tsv')\n",
    "        self.valid_examples = self.processor.get_dev_examples(TSV_PATH + 'ET_VALID.tsv')\n",
    "        self.test_examples = self.processor.get_test_examples(TSV_PATH + 'ET_TEST.tsv')\n",
    "\n",
    "        self.train_loader = self.processor.create_dataloader(\n",
    "            examples=self.train_examples,\n",
    "            tokenizer=self.tokenizer,\n",
    "            max_length=self.max_length,\n",
    "            shuffle=True,\n",
    "            batch_size=self.batch_size,\n",
    "            use_pickle=self.use_pickle,\n",
    "        )\n",
    "        self.valid_loader = self.processor.create_dataloader(\n",
    "            examples=self.valid_examples,\n",
    "            tokenizer=self.tokenizer,\n",
    "            max_length=self.max_length,\n",
    "            shuffle=False,\n",
    "            batch_size=self.batch_size,\n",
    "            use_pickle=self.use_pickle,\n",
    "        )\n",
    "        self.test_loader = self.processor.create_dataloader(\n",
    "            examples=self.test_examples,\n",
    "            tokenizer=self.tokenizer,\n",
    "            max_length=self.max_length,\n",
    "            shuffle=False,\n",
    "            batch_size=self.batch_size,\n",
    "            use_pickle=self.use_pickle,\n",
    "        )\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        input_ids, attention_mask, token_type_ids, labels = batch\n",
    "        outputs = self(input_ids, attention_mask, token_type_ids)\n",
    "        loss = self.criterion(outputs, labels)\n",
    "\n",
    "        _, preds = torch.max(outputs, dim=1)\n",
    "        acc = (preds == labels).float().mean()\n",
    "\n",
    "        tensorboard_logs = {'train_loss': loss, 'train_acc': acc}\n",
    "        return {'loss': loss, 'log': tensorboard_logs, 'progress_bar': tensorboard_logs}\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        input_ids, attention_mask, token_type_ids, labels = batch\n",
    "        outputs = self(input_ids, attention_mask, token_type_ids)\n",
    "        loss = self.criterion(outputs, labels)\n",
    "\n",
    "        _, preds = torch.max(outputs, dim=1)\n",
    "        acc = (preds == labels).float().mean()\n",
    "\n",
    "        return {'val_loss': loss, 'val_acc': acc}\n",
    "\n",
    "    def validation_epoch_end(self, outputs):\n",
    "        val_loss = torch.stack([x['val_loss'] for x in outputs]).mean()\n",
    "        val_acc = torch.stack([x['val_acc'] for x in outputs]).mean()\n",
    "\n",
    "        tensorboard_logs = {'val_loss': val_loss, 'val_acc': val_acc}\n",
    "        return {'val_loss': val_loss, 'log': tensorboard_logs, 'progress_bar': tensorboard_logs}\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        return torch.optim.Adam([p for p in self.parameters() if p.requires_grad], lr=2e-5, eps=1e-8)\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return self.train_loader\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        return self.valid_loader\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EntityTypingPredictor:\n",
    "\n",
    "    def __init__(self, ckpt_name, batch_size=8, use_pickle=True):\n",
    "        self.ckpt_name = ckpt_name\n",
    "        self.batch_size = batch_size\n",
    "        self.use_pickle = use_pickle\n",
    "\n",
    "    def generate_tsv_result(self, tsv_name, tsv_type='Valid'):\n",
    "        processor = EntityTypingProcessor()\n",
    "        tokenizer = BertTokenizer.from_pretrained(\"hfl/chinese-roberta-wwm-ext\")\n",
    "\n",
    "        if tsv_type == 'Valid':\n",
    "            examples = processor.get_dev_examples(TSV_PATH + tsv_name)\n",
    "        elif tsv_type == 'Test':\n",
    "            examples = processor.get_test_examples(TSV_PATH + tsv_name)\n",
    "        else:\n",
    "            raise ValueError('tsv_type error')\n",
    "        dataloader = processor.create_dataloader(\n",
    "            examples=examples,\n",
    "            tokenizer=tokenizer,\n",
    "            max_length=64,\n",
    "            shuffle=False,\n",
    "            batch_size=self.batch_size,\n",
    "            use_pickle=self.use_pickle,\n",
    "        )\n",
    "\n",
    "        model = EntityTypingModel.load_from_checkpoint(\n",
    "            checkpoint_path=CKPT_PATH + self.ckpt_name,\n",
    "        )\n",
    "        model.to(DEVICE)\n",
    "        model = nn.DataParallel(model)\n",
    "        model.eval()\n",
    "\n",
    "        result_list = []\n",
    "        for batch in tqdm(dataloader):\n",
    "            for i in range(len(batch)):\n",
    "                batch[i] = batch[i].to(DEVICE)\n",
    "\n",
    "            input_ids, attention_mask, token_type_ids, labels = batch\n",
    "            outputs = model(input_ids, attention_mask, token_type_ids)\n",
    "            _, preds = torch.max(outputs, dim=1)\n",
    "            result_list.extend(preds.tolist())\n",
    "\n",
    "        idx_to_type = PICKLE_DATA['IDX_TO_TYPE']\n",
    "        result_list = [idx_to_type[x] for x in result_list]\n",
    "        tsv_data = pd.read_csv(TSV_PATH + tsv_name, sep='\\t')\n",
    "        tsv_data['result'] = result_list\n",
    "        result_name = tsv_name.split('.')[0] + '_RESULT.tsv'\n",
    "        tsv_data.to_csv(RESULT_PATH + result_name, index=False, sep='\\t')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from collections import defaultdict\n",
    "\n",
    "\n",
    "class Eval(object):\n",
    "    \"\"\"\n",
    "    Entity Linking Evaluation\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, golden_file_path, user_file_path):\n",
    "        self.golden_file_path = golden_file_path\n",
    "        self.user_file_path = user_file_path\n",
    "        self.tp = 0\n",
    "        self.fp = 0\n",
    "        self.total_recall = 0\n",
    "        self.errno = None\n",
    "\n",
    "    def format_check(self, file_path):\n",
    "        \"\"\"\n",
    "        文件格式验证\n",
    "        :param file_path: 文件路径\n",
    "        :return: Bool类型：是否通过格式检查，通过为True，反之False\n",
    "        \"\"\"\n",
    "        flag = True\n",
    "        for line in open(file_path, encoding='utf-8'):\n",
    "            json_info = json.loads(line.strip())\n",
    "            if 'text_id' not in json_info:\n",
    "                flag = False\n",
    "                self.errno = 1\n",
    "                break\n",
    "            if 'text' not in json_info:\n",
    "                flag = False\n",
    "                self.errno = 2\n",
    "                break\n",
    "            if 'mention_data' not in json_info:\n",
    "                flag = False\n",
    "                self.errno = 3\n",
    "                break\n",
    "            if not json_info['text_id'].isdigit():\n",
    "                flag = False\n",
    "                self.errno = 5\n",
    "                break\n",
    "            if not isinstance(json_info['mention_data'], list):\n",
    "                flag = False\n",
    "                self.errno = 7\n",
    "                break\n",
    "            for mention_info in json_info['mention_data']:\n",
    "                if 'kb_id' not in mention_info:\n",
    "                    flag = False\n",
    "                    self.errno = 7\n",
    "                    break\n",
    "                if 'mention' not in mention_info:\n",
    "                    flag = False\n",
    "                    self.errno = 8\n",
    "                    break\n",
    "                if 'offset' not in mention_info:\n",
    "                    flag = False\n",
    "                    self.errno = 9\n",
    "                    break\n",
    "                if not mention_info['offset'].isdigit():\n",
    "                    flag = False\n",
    "                    self.errno = 13\n",
    "                    break\n",
    "        return flag\n",
    "\n",
    "    def micro_f1(self):\n",
    "        \"\"\"\n",
    "        :return: float类型：精确率，召回率，Micro-F1值\n",
    "        \"\"\"\n",
    "        # 文本格式验证\n",
    "        flag_golden = self.format_check(self.golden_file_path)\n",
    "        flag_user = self.format_check(self.user_file_path)\n",
    "        # 格式验证失败直接返回None\n",
    "        if not flag_golden or not flag_user:\n",
    "            return None, None, None\n",
    "        precision = 0\n",
    "        recall = 0\n",
    "        self.tp = 0\n",
    "        self.fp = 0\n",
    "        self.total_recall = 0\n",
    "        golden_dict = defaultdict(list)\n",
    "        for line in open(self.golden_file_path, encoding='utf-8'):\n",
    "            golden_info = json.loads(line.strip())\n",
    "            text_id = golden_info['text_id']\n",
    "            text = golden_info['text']\n",
    "            mention_data = golden_info['mention_data']\n",
    "            for mention_info in mention_data:\n",
    "                kb_id = mention_info['kb_id']\n",
    "                mention = mention_info['mention']\n",
    "                offset = mention_info['offset']\n",
    "                key = '\\1'.join([text_id, text, mention, offset]).encode('utf8')\n",
    "                # value的第二个元素表示标志位，用于判断是否已经进行了统计\n",
    "                golden_dict[key] = [kb_id, 0]\n",
    "                self.total_recall += 1\n",
    "\n",
    "        # 进行评估\n",
    "        for line in open(self.user_file_path, encoding='utf-8'):\n",
    "            golden_info = json.loads(line.strip())\n",
    "            text_id = golden_info['text_id']\n",
    "            text = golden_info['text']\n",
    "            mention_data = golden_info['mention_data']\n",
    "            for mention_info in mention_data:\n",
    "                kb_id = mention_info['kb_id']\n",
    "                mention = mention_info['mention']\n",
    "                offset = mention_info['offset']\n",
    "                key = '\\1'.join([text_id, text, mention, offset]).encode('utf8')\n",
    "                if key in golden_dict:\n",
    "                    kb_result_golden = golden_dict[key]\n",
    "                    if kb_id.isdigit():\n",
    "                        if kb_id in [kb_result_golden[0]] and kb_result_golden[1] in [0]:\n",
    "                            self.tp += 1\n",
    "                        else:\n",
    "                            self.fp += 1\n",
    "                    else:\n",
    "                        # nil golden结果\n",
    "                        nil_res = kb_result_golden[0].split('|')\n",
    "                        if kb_id in nil_res and kb_result_golden[1] in [0]:\n",
    "                            self.tp += 1\n",
    "                        else:\n",
    "                            self.fp += 1\n",
    "                    golden_dict[key][1] = 1\n",
    "                else:\n",
    "                    self.fp += 1\n",
    "        if self.tp + self.fp > 0:\n",
    "            precision = float(self.tp) / (self.tp + self.fp)\n",
    "        if self.total_recall > 0:\n",
    "            recall = float(self.tp) / self.total_recall\n",
    "        a = 2 * precision * recall\n",
    "        b = precision + recall\n",
    "        if b == 0:\n",
    "            return 0, 0, 0\n",
    "        f1 = a / b\n",
    "        return precision, recall, f1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_pickle_file():\n",
    "    print(\"preprocess_pickle_file .......\")\n",
    "    processor = PicklePreprocessor()\n",
    "    processor.run()\n",
    "    read_pickle()\n",
    "\n",
    "\n",
    "def preprocess_tsv_file():\n",
    "    print(\"preprocess_tsv_file .......\")\n",
    "    processor = DataFramePreprocessor()\n",
    "    processor.run()\n",
    "\n",
    "\n",
    "def generate_feature_pickle():\n",
    "    print(\"generate_feature_pickle .......\")\n",
    "    processor = EntityLinkingProcessor()\n",
    "    processor.generate_feature_pickle(max_length=384)\n",
    "\n",
    "    processor = EntityTypingProcessor()\n",
    "    processor.generate_feature_pickle(max_length=64)\n",
    "\n",
    "\n",
    "def train_entity_linking_model(ckpt_name):\n",
    "    print(\"train_entity_linking_model .......\")\n",
    "    model = EntityLinkingModel(max_length=384, batch_size=32)\n",
    "    trainer = pl.Trainer(\n",
    "        max_epochs=1,\n",
    "#         gpus=1,\n",
    "#         distributed_backend='dp',\n",
    "#         default_save_path=EL_SAVE_PATH,\n",
    "        profiler=True,\n",
    "    )\n",
    "    trainer.fit(model)\n",
    "    trainer.save_checkpoint(CKPT_PATH + ckpt_name)\n",
    "\n",
    "\n",
    "def train_entity_typing_model(ckpt_name):\n",
    "    print(\"train_entity_typing_model .......\")\n",
    "    model = EntityTypingModel(max_length=64, batch_size=64)\n",
    "    trainer = pl.Trainer(\n",
    "        max_epochs=1,\n",
    "#         gpus=1,\n",
    "#         distributed_backend='dp',\n",
    "#         default_save_path=ET_SAVE_PATH,\n",
    "        profiler=True,\n",
    "    )\n",
    "    trainer.fit(model)\n",
    "    trainer.save_checkpoint(CKPT_PATH + ckpt_name)\n",
    "\n",
    "\n",
    "def generate_link_tsv_result(ckpt_name):\n",
    "    print(\"generate_link_tsv_result .......\")\n",
    "    predictor = EntityLinkingPredictor(ckpt_name, batch_size=24, use_pickle=True)\n",
    "    predictor.generate_tsv_result('EL_VALID.tsv', tsv_type='Valid')\n",
    "    predictor.generate_tsv_result('EL_TEST.tsv', tsv_type='Test')\n",
    "\n",
    "\n",
    "def generate_type_tsv_result(ckpt_name):\n",
    "    print(\"generate_type_tsv_result .......\")\n",
    "    predictor = EntityTypingPredictor(ckpt_name, batch_size=64, use_pickle=True)\n",
    "    predictor.generate_tsv_result('ET_VALID.tsv', tsv_type='Valid')\n",
    "    predictor.generate_tsv_result('ET_TEST.tsv', tsv_type='Test')\n",
    "\n",
    "\n",
    "def make_predication_result(input_name, output_name, el_ret_name, et_ret_name):\n",
    "    print(\"make_predication_result .......\")\n",
    "    entity_to_kbids = PICKLE_DATA['ENTITY_TO_KBIDS']\n",
    "\n",
    "    el_ret = pd.read_csv(\n",
    "        RESULT_PATH + el_ret_name, sep='\\t', dtype={\n",
    "            'text_id': np.str_,\n",
    "            'offset': np.str_,\n",
    "            'kb_id': np.str_\n",
    "        })\n",
    "\n",
    "    et_ret = pd.read_csv(RESULT_PATH + et_ret_name, sep='\\t', dtype={'text_id': np.str_, 'offset': np.str_})\n",
    "\n",
    "    result = []\n",
    "    with open(RAW_PATH + input_name, 'r', encoding=\"utf-8\") as f:\n",
    "        for line in tqdm(f):\n",
    "            line = json.loads(line)\n",
    "            for data in line['mention_data']:\n",
    "                text_id = line['text_id']\n",
    "                offset = data['offset']\n",
    "\n",
    "                candidate_data = el_ret[(el_ret['text_id'] == text_id) & (el_ret['offset'] == offset)]\n",
    "                # Entity Linking\n",
    "                if len(candidate_data) > 0 and candidate_data['logits'].max() > 0:\n",
    "                    max_idx = candidate_data['logits'].idxmax()\n",
    "                    data['kb_id'] = candidate_data.loc[max_idx]['kb_id']\n",
    "                # Entity Typing\n",
    "                else:\n",
    "                    type_data = et_ret[(et_ret['text_id'] == text_id) & (et_ret['offset'] == offset)]\n",
    "                    data['kb_id'] = 'NIL_' + type_data.iloc[0]['result']\n",
    "            result.append(line)\n",
    "\n",
    "    with open(RESULT_PATH + output_name, 'w', encoding=\"utf-8\") as f:\n",
    "        for r in result:\n",
    "            json.dump(r, f, ensure_ascii=False)\n",
    "            f.write('\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "6953it [00:00, 36015.38it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "preprocess_pickle_file .......\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "324418it [00:08, 40400.13it/s]\n",
      "7170it [00:00, 34816.90it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "preprocess_tsv_file .......\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "70000it [00:01, 35652.74it/s]\n",
      "10000it [00:00, 31778.33it/s]\n",
      "10000it [00:00, 17845.10it/s]\n",
      "70000it [00:00, 114552.95it/s]\n",
      "10000it [00:00, 106838.66it/s]\n",
      "10000it [00:00, 160964.03it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "generate_feature_pickle .......\n"
     ]
    }
   ],
   "source": [
    "set_random_seed(20200707)\n",
    "preprocess_pickle_file()  # 生成全局变量索引pickle文件\n",
    "preprocess_tsv_file()  # 生成模型训练、验证、推断所需的tsv文件\n",
    "generate_feature_pickle()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_entity_linking_model .......\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: False, used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "\n",
      "  | Name      | Type                          | Params\n",
      "------------------------------------------------------------\n",
      "0 | bert      | BertForSequenceClassification | 102 M \n",
      "1 | criterion | BCEWithLogitsLoss             | 0     \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "finish\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/vodka/anaconda3/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 12 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  warnings.warn(*args, **kwargs)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Validation sanity check', layout=Layout…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\r"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/vodka/anaconda3/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: The dataloader, train dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 12 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.\n",
      "  warnings.warn(*args, **kwargs)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9e42be004d814ae18fae31c113afed84",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=1.0, bar_style='info', description='Training', layout=Layout(flex='2'), max…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/vodka/anaconda3/lib/python3.7/site-packages/pytorch_lightning/utilities/distributed.py:25: UserWarning: Detected KeyboardInterrupt, attempting graceful shutdown...\n",
      "  warnings.warn(*args, **kwargs)\n",
      "\n",
      "\n",
      "Profiler Report\n",
      "\n",
      "Action              \t|  Mean duration (s)\t|  Total time (s) \n",
      "-----------------------------------------------------------------\n",
      "on_train_start      \t|  0.01388        \t|  0.01388        \n",
      "on_epoch_start      \t|  0.0015656      \t|  0.0015656      \n",
      "get_train_batch     \t|  0.088284       \t|  0.35314        \n",
      "on_batch_start      \t|  2.8685e-05     \t|  0.00011474     \n",
      "model_forward       \t|  8.7427         \t|  34.971         \n",
      "model_backward      \t|  13.69          \t|  54.76          \n",
      "on_after_backward   \t|  1.9803e-06     \t|  5.941e-06      \n",
      "optimizer_step      \t|  0.38785        \t|  1.1636         \n",
      "on_batch_end        \t|  0.025867       \t|  0.077601       \n",
      "on_train_end        \t|  0.0015541      \t|  0.0015541      \n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] No such file or directory: './ckpt/EL_BASE_EPOCH0.ckpt.part'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-20-1520292c1bf6>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mtrain_entity_linking_model\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'EL_BASE_EPOCH0.ckpt'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      2\u001b[0m \u001b[0mgenerate_link_tsv_result\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'EL_BASE_EPOCH0.ckpt'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[0mel_ret\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mread_csv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"./data/result/ET_TEST_RESULT.tsv\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msep\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'\\t'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mel_ret\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mhead\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m<ipython-input-19-4fa761044230>\u001b[0m in \u001b[0;36mtrain_entity_linking_model\u001b[0;34m(ckpt_name)\u001b[0m\n\u001b[1;32m     32\u001b[0m     )\n\u001b[1;32m     33\u001b[0m     \u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 34\u001b[0;31m     \u001b[0mtrainer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msave_checkpoint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mCKPT_PATH\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0mckpt_name\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     35\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     36\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/training_io.py\u001b[0m in \u001b[0;36msave_checkpoint\u001b[0;34m(self, filepath, weights_only)\u001b[0m\n\u001b[1;32m    271\u001b[0m             \u001b[0;31m# do the actual save\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    272\u001b[0m             \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 273\u001b[0;31m                 \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_atomic_save\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcheckpoint\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfilepath\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    274\u001b[0m             \u001b[0;32mexcept\u001b[0m \u001b[0mAttributeError\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0merr\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    275\u001b[0m                 \u001b[0;32mif\u001b[0m \u001b[0mLightningModule\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mCHECKPOINT_HYPER_PARAMS_KEY\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mcheckpoint\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/pytorch_lightning/trainer/training_io.py\u001b[0m in \u001b[0;36m_atomic_save\u001b[0;34m(self, checkpoint, filepath)\u001b[0m\n\u001b[1;32m    262\u001b[0m         \"\"\"\n\u001b[1;32m    263\u001b[0m         \u001b[0mtmp_path\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfilepath\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m+\u001b[0m \u001b[0;34m\".part\"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 264\u001b[0;31m         \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msave\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcheckpoint\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtmp_path\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    265\u001b[0m         \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreplace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtmp_path\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfilepath\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    266\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/torch/serialization.py\u001b[0m in \u001b[0;36msave\u001b[0;34m(obj, f, pickle_module, pickle_protocol)\u001b[0m\n\u001b[1;32m    258\u001b[0m         \u001b[0;34m>>\u001b[0m\u001b[0;34m>\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msave\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbuffer\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    259\u001b[0m     \"\"\"\n\u001b[0;32m--> 260\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0m_with_file_like\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"wb\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mlambda\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0m_save\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobj\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpickle_module\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpickle_protocol\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    261\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    262\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/torch/serialization.py\u001b[0m in \u001b[0;36m_with_file_like\u001b[0;34m(f, mode, body)\u001b[0m\n\u001b[1;32m    181\u001b[0m             \u001b[0;34m(\u001b[0m\u001b[0msys\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mversion_info\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;36m3\u001b[0m \u001b[0;32mand\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpathlib\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mPath\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    182\u001b[0m         \u001b[0mnew_fd\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mTrue\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 183\u001b[0;31m         \u001b[0mf\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmode\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    184\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    185\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mbody\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: './ckpt/EL_BASE_EPOCH0.ckpt.part'"
     ]
    }
   ],
   "source": [
    "train_entity_linking_model('EL_BASE_EPOCH0.ckpt')\n",
    "generate_link_tsv_result('EL_BASE_EPOCH0.ckpt')\n",
    "el_ret = pd.read_csv(\"./data/result/ET_TEST_RESULT.tsv\", sep='\\t')\n",
    "print(el_ret.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
